#include "StackAllocator.h"
#include "Atomic.h"
#include "Assertions.h"

#include <algorithm>

namespace CEngine
{

StackAllocator::StackAllocator(U64 stackSize)
    :   freeBytes(stackSize),
        totalBytes(stackSize)
{
    stackTop = stackPointer = new U8[stackSize];
}

StackAllocator::StackAllocator(const StackAllocator& stack)
    :   freeBytes(stack.totalBytes),
        totalBytes(stack.totalBytes)
{
    stackPointer = new U8[totalBytes];

    std::copy((U8*)stack.stackPointer, (U8*)stack.stackPointer + stack.totalBytes, (U8*)stackPointer);
    stackTop = (U8*)stackPointer + ((U8*)stack.stackTop - (U8*)stack.stackPointer);
}

StackAllocator::StackAllocator(const StackAllocator& stack, U64 stackSize)
    :   freeBytes(stackSize),
        totalBytes(stackSize)
{
    ASSERT(stackSize >= stack.totalBytes);

    stackPointer = new U8[stackSize];

    std::copy((U8*)stack.stackPointer, (U8*)stack.stackPointer + stack.totalBytes, (U8*)stackPointer);
    stackTop = (U8*)stackPointer + ((U8*)stack.stackTop - (U8*)stack.stackPointer);
}

StackAllocator::~StackAllocator(void)
{
    delete[] (U8*)stackPointer;
    stackPointer = NULL;
    stackTop = NULL;
    freeBytes = 0;
    totalBytes = 0;
}

void*
StackAllocator::alloc(U64 bytes) {
    if (bytes > freeBytes)
        return NULL;

    void* temp = stackTop;

    stackTop += bytes;
    freeBytes -= bytes;

    return temp;
}

void*
StackAllocator::alignedAlloc(U64 bytes, U64 alignment) {
    ///use normal alloc if alignment == 1
    ASSERT(alignment > 1);

    Marker marker = getMarker();

    U64 expandedBytes = bytes + alignment;

    U64 rawAddress = (U64)alloc(expandedBytes);
    if (rawAddress == NULL)
        return NULL;

    U64 mask = alignment - 1;
    U64 misalignment = (rawAddress & mask);
    U64 adjustment = alignment - misalignment;
    U64 alignedAddress = rawAddress + adjustment;

    freeToMarker((marker + expandedBytes) - adjustment);

    return (void*)alignedAddress;
}

void
StackAllocator::freeToMarker(Marker marker) {
    if ((void*)marker == NULL) {
        return;
    }
    freeBytes += (U64)stackTop - marker;
    stackTop = (void*)marker;
}

void
StackAllocator::clear(void) {
    stackTop = stackPointer;
    freeBytes = totalBytes;
}

}
